#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelGatherV3_32 {
public:
    __aicore__ inline KernelGatherV3_32() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indices, GM_ADDR axis, GM_ADDR out, 
                                int32_t batchs, int32_t repeat, int32_t cycles, int32_t interval, 
                                int32_t batch_lenght, int32_t repeat_lenght, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->batchs = batchs;
        this->repeat = repeat;
        this->cycles = cycles;
        this->interval = interval;
        this->batch_lenght = batch_lenght;
        this->repeat_lenght = repeat_lenght;
        this->tileDataLenght = 126976/sizeof(DTYPE_X);
        this->loopCount = this->cycles / this->tileDataLenght;
        this->tailDataLenght = this->cycles - this->loopCount * this->tileDataLenght;
        if(this->tailDataLenght == 0)
        {
            this->loopCount -= 1;
            this->tailDataLenght = this->tileDataLenght;
        }
        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x));
        indicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(indices));
        outGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(out));
        pipeIn->InitBuffer(inOutQueue, BUFFER_NUM, this->interval * this->processDataLenght * sizeof(DTYPE_X));
    }
    __aicore__ inline void Process() {
        int32_t num = 0;
        int32_t pian1 = 0;
        int32_t pian_indices = 0;
        int32_t pian2 = 0;
        for(int32_t b=0; b < this->batchs; b++)
        {
            pian2 = pian1;
            for(int32_t i = 0; i < this->repeat; i++)
            {
                this->processDataLenght = this->tileDataLenght;
                for(int32_t j = 0; j < this->loopCount; j++)
                {
                    CopyIn(pian_indices + j*this->tileDataLenght, pian2);
                    CopyOut(num);
                    num += this->interval * this->processDataLenght;
                }
                this->processDataLenght = this->tailDataLenght;
                CopyIn(pian_indices + this->loopCount*this->tileDataLenght, pian2);
                CopyOut(num);
                num += this->interval * this->processDataLenght;
                pian2 += repeat_lenght;
            }
            pian1 += batch_lenght;
            pian_indices += this->cycles;
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t offset1, int32_t pian2) {
        LocalTensor<DTYPE_X> bindLocal = inOutQueue.AllocTensor<DTYPE_X>();
        for(int i=0; i<this->processDataLenght; i++)
        {
            DataCopy(bindLocal[this->interval * i], xGm[pian2 +indicesGm.GetValue(offset1 + i)* this->interval], this->interval);
        }
        inOutQueue.EnQue(bindLocal);
    }
    __aicore__ inline void CopyOut(int32_t offset) {
        LocalTensor<DTYPE_X> bindLocal = inOutQueue.DeQue<DTYPE_X>();
        DataCopy(outGm[offset], bindLocal, this->interval*this->processDataLenght);
        inOutQueue.FreeTensor(bindLocal);
    }
private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<int32_t> indicesGm;
    GlobalTensor<DTYPE_X> outGm;
    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> inOutQueue; 
    int32_t batchs;
    int32_t repeat;
    int32_t cycles;
    int32_t interval;
    int32_t batch_lenght;
    int32_t repeat_lenght;
    int32_t loopCount;
    int32_t tileDataLenght;
    int32_t tailDataLenght;
    int32_t processDataLenght;
};
class KernelGatherV3 {
public:
    __aicore__ inline KernelGatherV3() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indices, GM_ADDR axis, GM_ADDR out, 
                                int32_t batchs, int32_t repeat, int32_t cycles, int32_t interval, 
                                int32_t batch_lenght, int32_t repeat_lenght, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->batchs = batchs;
        this->repeat = repeat;
        this->cycles = cycles;
        this->interval = interval;
        this->batch_lenght = batch_lenght;
        this->repeat_lenght = repeat_lenght;
        this->interval_32 = (this->interval*sizeof(DTYPE_X)+31)/32*32/sizeof(DTYPE_X);
        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x));
        indicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(indices));
        outGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(out));
        pipeIn->InitBuffer(inOutQueue, BUFFER_NUM, this->interval_32 * sizeof(DTYPE_X));
    }
    __aicore__ inline void Process() {
        int32_t num = 0;
        for(int32_t b=0; b < this->batchs; b++)
        {
            int32_t pian1 = b*batch_lenght;
            int32_t pian_indices = b*this->cycles;
            for(int32_t i = 0; i < this->repeat; i++)
            {
                int32_t pian2 = pian1 + i*repeat_lenght;
                for(int32_t j = 0; j < this->cycles; j++)
                {
                    int32_t indices = indicesGm.GetValue(pian_indices + j);
                    int32_t pian = pian2 + indices * this->interval;
                    CopyIn(pian);
                    CopyOut(num);
                    num += this->interval;
                }
            }
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t offset) {
        LocalTensor<DTYPE_X> bindLocal = inOutQueue.AllocTensor<DTYPE_X>();
        DataCopy(bindLocal, xGm[offset], this->interval_32);
        inOutQueue.EnQue(bindLocal);
    }
    __aicore__ inline void CopyOut(int32_t offset) {
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(this->interval * sizeof(DTYPE_X)), 0, 0, 0};
        LocalTensor<DTYPE_X> bindLocal = inOutQueue.DeQue<DTYPE_X>();
        DataCopyPad(outGm[offset], bindLocal, copyParams);
        inOutQueue.FreeTensor(bindLocal);
    }
private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<int32_t> indicesGm;
    GlobalTensor<DTYPE_X> outGm;
    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> inOutQueue; 
    int32_t batchs;
    int32_t repeat;
    int32_t cycles;
    int32_t interval;
    int32_t interval_32;
    int32_t batch_lenght;
    int32_t repeat_lenght;
};
class KernelGatherV3_1 {
public:
    __aicore__ inline KernelGatherV3_1() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR indices, GM_ADDR axis, GM_ADDR out, 
                                int32_t batchs, int32_t repeat, int32_t cycles, int32_t interval, 
                                int32_t batch_lenght, int32_t repeat_lenght, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->batchs = batchs;
        this->repeat = repeat;
        this->cycles = cycles;
        this->interval = interval;
        this->batch_lenght = batch_lenght;
        this->repeat_lenght = repeat_lenght;
        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x));
        indicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(indices));
        outGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(out));
    }
    __aicore__ inline void Process() {
        int32_t num = 0;
        for(int32_t b=0; b < this->batchs; b++)
        {
            int32_t pian1 = b*batch_lenght;
            int32_t pian_indices = b*this->cycles;
            for(int32_t i = 0; i < this->repeat; i++)
            {
                int32_t pian2 = pian1 + i*repeat_lenght;
                for(int32_t j = 0; j < this->cycles; j++)
                {
                    int32_t indices = indicesGm.GetValue(pian_indices + j);
                    int32_t pian = pian2 + indices * this->interval;
                    for(int32_t z = 0; z < this->interval; z++)
                    {
                        outGm.SetValue(num, xGm.GetValue(pian));
                        num ++;
                        pian ++;
                    }
                }
            }
        }
    }
private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<int32_t> indicesGm;
    GlobalTensor<DTYPE_X> outGm;
    int32_t batchs;
    int32_t repeat;
    int32_t cycles;
    int32_t interval;
    int32_t batch_lenght;
    int32_t repeat_lenght;
};
extern "C" __global__ __aicore__ void gather_v3(GM_ADDR x, GM_ADDR indices, GM_ADDR axis, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    if (TILING_KEY_IS(1)) {
        KernelGatherV3_32 op;
        op.Init(x, indices, axis, out,
                tiling_data.batchs, tiling_data.repeat, tiling_data.cycles, tiling_data.interval, 
                tiling_data.batch_lenght, tiling_data.repeat_lenght, &pipe);
        op.Process();
    }
    else if (TILING_KEY_IS(2)) {
        KernelGatherV3 op;
        op.Init(x, indices, axis, out,
                tiling_data.batchs, tiling_data.repeat, tiling_data.cycles, tiling_data.interval, 
                tiling_data.batch_lenght, tiling_data.repeat_lenght, &pipe);
        op.Process();
    }
    else if (TILING_KEY_IS(3)) {
        KernelGatherV3_1 op;
        op.Init(x, indices, axis, out,
                tiling_data.batchs, tiling_data.repeat, tiling_data.cycles, tiling_data.interval, 
                tiling_data.batch_lenght, tiling_data.repeat_lenght, &pipe);
        op.Process();
    }
}